import matplotlib.pyplot as plt
import os
from os.path import join, dirname, exists, split
from os import listdir
from collections import defaultdict
from glob import glob
import pandas as pd
import mbrl.util.common

exp_path = join(dirname(dirname(dirname(__file__))), "exp")


def try_read_csv(path, shortest_len=50):
    if not exists(path):
        return None
    try:
        df = pd.read_csv(path)
    except pd.errors.EmptyDataError:
        return None
    if len(df) < shortest_len:
        return None
    return df


def draw(env_name, experiment_names=None, shortest_len=5):
    if experiment_names is None:
        experiment_names = ["default"]

    env_dir_name = '{}'.format(env_name)
    episode_reward_dict = defaultdict(lambda: defaultdict(list))
    env_step = []
    for algo_name in listdir(exp_path):
        for experiment_name in experiment_names:
            seeds = []
            lengths = []
            p = join(exp_path, algo_name, experiment_name, env_dir_name)
            if exists(p):
                for single_exp_path in glob(join(p, "*", "*")):
                    df = try_read_csv(join(single_exp_path, "results.csv"), shortest_len)
                    if df is None:
                        print(single_exp_path)
                        continue

                    cfg = mbrl.util.common.load_hydra_cfg(single_exp_path)
                    seeds.append(cfg.seed)
                    lengths.append(len(df))
                    episode_reward_dict[algo_name][experiment_name].append(df["episode_reward"])
                    if len(df) > len(env_step):
                        env_step = df["env_step"]
            if len(seeds) != 0:
                print("{}:{}".format(algo_name.upper(), experiment_name), (seeds), (lengths))

    fig = plt.figure()
    for algo_name in episode_reward_dict:
        for experiment_name in episode_reward_dict[algo_name]:
            df = pd.DataFrame()
            print(algo_name, len(episode_reward_dict[algo_name][experiment_name]))
            for i, value in enumerate(episode_reward_dict[algo_name][experiment_name]):
                df[i] = value
            mean = df.mean(1)
            std_down = mean - df.std(1) * 1
            std_up = mean + df.std(1) * 1
            plt.plot(env_step[:len(mean)] / 1000, mean, label="{}:{}".format(algo_name.upper(), experiment_name))
            plt.fill_between(env_step[:len(mean)] / 1000, std_down, std_up, alpha=0.2)
    plt.legend()
    plt.show()

    pass


if __name__ == '__main__':
    draw("gym___Hopper-v2", ["discovery-time_scale-1-long-rollout"])
